import os, sys
print(os.getcwd())
sys.path.append(os.getcwd())
from skorch.core import *
from skorch.vision import *
import nrrd
import PIL.Image
import PIL.ImageDraw


def brightness_status(mean_dev, mean):
    """
    return brightness status
    :param mean_dev:
    :param mean:
    :return:
    """
    if mean_dev < abs(mean):
        if mean > 0:
            return 1
        elif mean < 0:
            return 2
        else:
            return 0
    else:
        return 0


def adjust_brightness(dcm_info, refer:int=128):
    """
    adjust dicom image brightness
    :param image_path:
    :param refer:default 128
    :return:
    """

    def image_mean(image: NPImage, refer: int = 128):
        """
        compute mean_dev, mean
        :param image:
        :param refer:
        :return:
        """
        mean = cv2.mean(image)[0]
        mean = mean - refer
        md = np.abs(image - refer - mean)
        mean_dev = np.mean(md)
        return mean_dev, mean

    default_wc = dcm_info['wc']
    image = convert_dcm_grayscale8(dcm_info)
    if image.max() == image.min():
        return False
    mean_dev, mean = image_mean(image, refer)
    status = brightness_status(mean_dev, mean)

    if status == 1:
        #print(image_path.name, ' 过亮 ', mean, mean_dev)
        dcm_info['wc'] = default_wc*0.8
        image = convert_dcm_grayscale8(dcm_info)
        mean_dev, mean = image_mean(image, refer)
        status = brightness_status(mean_dev, mean)
        if status == 1:
            dcm_info['wc'] = default_wc * 1.1
    return True


def label_colormap(N=256):

    def bitget(byteval, idx):
        return ((byteval & (1 << idx)) != 0)

    cmap = np.zeros((N, 3))
    for i in range(0, N):
        id = i
        r, g, b = 0, 0, 0
        for j in range(0, 8):
            r = np.bitwise_or(r, (bitget(id, 0) << 7 - j))
            g = np.bitwise_or(g, (bitget(id, 1) << 7 - j))
            b = np.bitwise_or(b, (bitget(id, 2) << 7 - j))
            id = (id >> 3)
        cmap[i, 0] = r
        cmap[i, 1] = g
        cmap[i, 2] = b
    cmap = cmap.astype(np.float32) / 255
    return cmap


def _validate_colormap(colormap, n_labels):
    if colormap is None:
        colormap = label_colormap(n_labels)
    else:
        assert colormap.shape == (colormap.shape[0], 3), \
            'colormap must be sequence of RGB values'
        assert 0 <= colormap.min() and colormap.max() <= 1, \
            'colormap must ranges 0 to 1'
    return colormap


# similar function as skimage.color.label2rgb
def label2rgb(
    lbl, img=None, n_labels=None, alpha=0.5, thresh_suppress=0, colormap=None,
):
    if n_labels is None:
        n_labels = len(np.unique(lbl))

    colormap = _validate_colormap(colormap, n_labels)
    colormap = (colormap * 255).astype(np.uint8)

    lbl_viz = colormap[lbl]
    lbl_viz[lbl == -1] = (0, 0, 0)  # unlabeled

    if img is not None:
        img_gray = PIL.Image.fromarray(img).convert('LA')
        img_gray = np.asarray(img_gray.convert('RGB'))
        # img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        # img_gray = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2RGB)
        lbl_viz = alpha * lbl_viz + (1 - alpha) * img_gray
        lbl_viz = lbl_viz.astype(np.uint8)

    return lbl_viz


def polygons_to_mask(img_shape, polygons):
    mask = np.zeros(img_shape[:2], dtype=np.uint8)
    mask = PIL.Image.fromarray(mask)
    xy = list(map(tuple, polygons))
    PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
    mask = np.array(mask, dtype=bool)
    return mask


def shapes_to_label(img_shape, shapes, label_name_to_value, type='class'):
    assert type in ['class', 'instance']

    cls = np.zeros([len(label_name_to_value)-1, img_shape[:2][0], img_shape[:2][1]], dtype=np.int32)
    if type == 'instance':
        ins = np.zeros([len(label_name_to_value)-1, img_shape[:2]], dtype=np.int32)
        instance_names = ['_background_']
    for shape in shapes:
        polygons = shape['points']
        label = shape['label']
        if type == 'class':
            cls_name = label.split('-')[0]
        elif type == 'instance':
            cls_name = label.split('-')[0]
            if label not in instance_names:
                instance_names.append(label)
            ins_id = len(instance_names) - 1
        cls_id = label_name_to_value[cls_name]
        mask = polygons_to_mask(img_shape[:2], polygons)
        cls[cls_id-1][mask] = cls_id
        if type == 'instance':
            ins[mask] = ins_id

    if type == 'instance':
        return cls, ins
    return cls


def lblsave(filename:Path, lbl:dict, class_dict:dict):
    if os.path.splitext(filename)[1] != '.png':
        filename += '.png'
    # Assume label ranses [-1, 254] for int32,
    # and [0, 255] for uint8 as VOC.
    if lbl.min() >= -1 and lbl.max() < 255:
        lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode='P')
        colormap = label_colormap(255)
        lbl_pil.putpalette((colormap * 255).astype(np.uint8).flatten())
        lbl_pil.save(filename)
    else:
        logging.warning(
            '[%s] Cannot save the pixel-wise class label as PNG, '
            'so please use the npy file.' % filename
        )


def lbl_save_multi_binary(filename:Path, lbl, class_dict:dict):
    for i in range(lbl.shape[0]):
        mask_key = list(class_dict.keys())[list(class_dict.values()).index(i+1)]
        cur_filename = filename.parents[0].joinpath(mask_key)
        if not cur_filename.exists():
            os.makedirs(str(cur_filename))
        cur_filename = cur_filename.joinpath(filename.name)
        if cur_filename.exists():
            continue
        cur_lbl = lbl[i]
        if cur_lbl.min() >= -1 and cur_lbl.max() < 255:
            cur_lbl_idx = cur_lbl.max()
            cur_lbl[cur_lbl<cur_lbl_idx]=0
            cur_lbl = cur_lbl/cur_lbl_idx*255
            cur_lbl_out = cur_lbl.transpose((1,0))[..., np.newaxis]
            if mask_key == 'lung':
                nrrd.write(str(filename), cur_lbl_out.astype(np.uint8))
            # lbl_pil = PIL.Image.fromarray(cur_lbl.astype(np.uint8), mode='L')
            # lbl_pil.save(os.path.splitext(cur_filename)[0] + '.png')
        else:
            logging.warning(
                '[%s] Cannot save the pixel-wise class label as PNG, '
                'so please use the npy file.' % filename
            )


def draw_label(label, img=None, label_names=None, colormap=None, **kwargs):
    """Draw pixel-wise label with colorization and label names.

    label: ndarray, (H, W)
        Pixel-wise labels to colorize.
    img: ndarray, (H, W, 3), optional
        Image on which the colorized label will be drawn.
    label_names: iterable
        List of label names.
    """
    import matplotlib.pyplot as plt

    backend_org = plt.rcParams['backend']
    plt.switch_backend('agg')

    plt.subplots_adjust(left=0, right=1, top=1, bottom=0,
                        wspace=0, hspace=0)
    plt.margins(0, 0)
    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    plt.gca().yaxis.set_major_locator(plt.NullLocator())

    if label_names is None:
        label_names = [str(l) for l in range(label.max() + 1)]

    colormap = _validate_colormap(colormap, len(label_names))

    label_viz = label2rgb(
        label, img, n_labels=len(label_names), colormap=colormap, **kwargs
    )
    plt.imshow(label_viz)
    plt.axis('off')

    plt_handlers = []
    plt_titles = []
    for label_value, label_name in enumerate(label_names):
        if label_value not in label:
            continue
        fc = colormap[label_value]
        p = plt.Rectangle((0, 0), 1, 1, fc=fc)
        plt_handlers.append(p)
        plt_titles.append('{value}: {name}'
                          .format(value=label_value, name=label_name))
    plt.legend(plt_handlers, plt_titles, loc='lower right', framealpha=.5)

    f = io.BytesIO()
    plt.savefig(f, bbox_inches='tight', pad_inches=0)
    plt.cla()
    plt.close()

    plt.switch_backend(backend_org)

    out_size = (label_viz.shape[1], label_viz.shape[0])
    out = PIL.Image.open(f).resize(out_size, PIL.Image.BILINEAR).convert('RGB')
    out = np.asarray(out)
    return out


def read_image(file_name):
    img = load_image_obj(file_name, 'pydcm')
    return img


def save_binary_masks(json_file, class_name_to_id, mask_dir):
    with open(json_file) as f:
        data = json.load(f)
        dcm_file = json_file.with_suffix('.dcm')
        image = read_image(str(dcm_file))

        if image is not None:
            cls = shapes_to_label(
                img_shape=image.shape,
                shapes=data['shapes'],
                label_name_to_value=class_name_to_id,
                type='class',
            )

            # out_png_file = mask_dir.joinpath(json_file.with_suffix('.png').name)
            out_png_file = mask_dir.joinpath(json_file.with_suffix('.nrrd').name)

            lbl_save_multi_binary(out_png_file, cls, class_name_to_id)
        else:
            print("file error ", str(dcm_file))


def save_dcn2npy(dcm_file:Path, save_dir:Path):
    dcm_dict = read_dcm_info(str(dcm_file))
    # check image brightness
    adjust_brightness(dcm_dict)
    # dicom convert gray image
    gray_image = convert_dcm_grayscale8(dcm_dict)
    npy_name = dcm_file.with_suffix(".png")
    if not save_dir.exists():
        save_dir.mkdir()
    npy_name = save_dir.joinpath(npy_name.name)

    # np.save(str(npy_name), gray_image)
    cv2.imwrite(str(npy_name), gray_image)


def load_json_flag(json_file, flag_text:Collection[str]):
    with open(str(json_file)) as f:
        flag_indexs = []
        data = json.load(f)
        for key in data['flags']:
            if key not in flag_text:
                print(json_file)
                break
        flags = data['flags']
        if flags['__ignore__']:
            return flag_indexs

        for key, value in flags.items():
            if value:
                flag_indexs.append(flag_text.index(key)-1)
        return flag_indexs


def process_dr_josn_flag():
    image_dir = Path('/home/blake/data/dataset/datasets/DR/dicom/DR/images')
    mask_dir = Path('/home/blake/data/dataset/datasets/DR/dicom/DR/masks')
    flag_text = ["__ignore__",
                 "excellent",
                 "FB_inside_body",
                 "FB_inside_lung_field",
                 "FB_outside_lung_field"]

    in_file_names = sorted(image_dir.glob('*.json'))

    labels_one_hot = []
    for json_file in tqdm(in_file_names):
        flag_indexs = load_json_flag(json_file, flag_text)
        if len(flag_indexs)==0:
            print(json_file)
            break
        hot_label = one_hot(flag_indexs, len(flag_text)-1)
        labels_one_hot.append(hot_label)
        # print(hot_label)

    labels_one_hot = np.asarray(labels_one_hot)

    print(labels_one_hot[:,0].sum(), labels_one_hot[:,1].sum(), labels_one_hot[:,2].sum(), labels_one_hot[:,3].sum())


def process_dr_diocm_josn(pool_flag=False):
    image_dir = Path('/home/cao/disk1/DR_dicoms/nrrd/images')
    mask_dir = Path('/home/cao/disk1/DR_dicoms/nrrd')

    class_names = ('__ignore__', 'lung', 'clavicle', 'shoulderblade')
    class_name_to_id = {'__ignore__': 0,
                        'lung': 1,
                        'clavicle': 2,
                        'shoulderblade': 3}

    colormap = label_colormap(255)

    in_file_names = sorted(image_dir.glob('*.json'))

    if pool_flag:
        pool = Pool(8)
        pool.map(partial(save_binary_masks, class_name_to_id=class_name_to_id, mask_dir=mask_dir), in_file_names)
        pool.close()
        pool.join()
    else:
        for json_file in tqdm(in_file_names):
            # json_file = Path('/home/blake/data/dataset/datasets/DR/dicom/DR2/images/
            # 0520181001343493-0000000003-0001-10001-1.2.156.147522.44.410947.0520181001343493.3.1.20181001103729.json')
            save_binary_masks(json_file, class_name_to_id=class_name_to_id, mask_dir=mask_dir)


def process_dr_diocm_npy(pool_flag=False):
    image_dir = Path('/home/blake/data/dataset/datasets/DR/dicom/DR2/')

    dcm_dir = image_dir.joinpath('dcms')
    in_file_names = sorted(dcm_dir.glob('*.dcm'))
    save_npy_dir = image_dir.joinpath('images')

    if pool_flag:
        pool = Pool(8)
        pool.map(partial(save_dcn2npy, save_dir=save_npy_dir), in_file_names)
        pool.close()
        pool.join()
    else:
        for json_file in tqdm(in_file_names):
            # json_file = Path('/home/blake/data/dataset/datasets/DR/dicom/DR2/images/
            # 0520181001343493-0000000003-0001-10001-1.2.156.147522.44.410947.0520181001343493.3.1.20181001103729.json')
            save_dcn2npy(json_file, save_dir=save_npy_dir)


if __name__ == '__main__':
    print('')
    # process_dr_josn_flag()
    process_dr_diocm_josn(pool_flag=False)
    # process_dr_diocm_npy(pool_flag=True)